import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import numpy as np
from PIL import Image
import torchvision.transforms as transforms

class PiDiNet(nn.Module):
    """PiDiNet 边缘检测网络"""
    
    def __init__(self, inplane=60, pdcs=None, dil=None, sa=False, rd=False):
        super(PiDiNet, self).__init__()
        self.sa = sa
        if pdcs is None:
            pdcs = [14, 14, 14, 14]
        if dil is None:
            dil = [1, 1, 1, 1]
        
        # 编码器部分
        self.conv1 = nn.Conv2d(3, inplane, kernel_size=7, stride=2, padding=3, bias=False)
        self.bn1 = nn.BatchNorm2d(inplane)
        self.relu = nn.ReLU(inplace=True)
        
        # 层级特征提取
        self.layer1 = self._make_layer(inplane, inplane, pdcs[0], dil[0], stride=1)
        self.layer2 = self._make_layer(inplane, inplane*2, pdcs[1], dil[1], stride=2)
        self.layer3 = self._make_layer(inplane*2, inplane*4, pdcs[2], dil[2], stride=2)
        self.layer4 = self._make_layer(inplane*4, inplane*8, pdcs[3], dil[3], stride=2)
        
        # 解码器和边缘预测
        self.decoder = nn.Sequential(
            nn.Conv2d(inplane*8, inplane*4, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(inplane*4, inplane*2, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(inplane*2, inplane, 3, padding=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(inplane, 1, 1)
        )
        
    def _make_layer(self, inplanes, planes, blocks, dilation, stride=1):
        layers = []
        for i in range(blocks):
            if i == 0:
                layers.append(nn.Conv2d(inplanes, planes, 3, stride=stride, 
                                      padding=dilation, dilation=dilation, bias=False))
            else:
                layers.append(nn.Conv2d(planes, planes, 3, padding=dilation, 
                                      dilation=dilation, bias=False))
            layers.append(nn.BatchNorm2d(planes))
            layers.append(nn.ReLU(inplace=True))
        return nn.Sequential(*layers)
    
    def forward(self, x):
        # 特征提取
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.relu(x)
        
        x1 = self.layer1(x)
        x2 = self.layer2(x1)
        x3 = self.layer3(x2)
        x4 = self.layer4(x3)
        
        # 解码和上采样
        out = self.decoder(x4)
        out = F.interpolate(out, size=x.shape[2:], mode='bilinear', align_corners=False)
        
        return torch.sigmoid(out)

class DocumentEdgeDetector:
    """文档边缘检测器"""
    
    def __init__(self, model_path, device='cpu'):
        self.device = device
        self.model = PiDiNet()
        
        # 加载预训练权重
        checkpoint = torch.load(model_path, map_location=device)
        if 'state_dict' in checkpoint:
            self.model.load_state_dict(checkpoint['state_dict'])
        else:
            self.model.load_state_dict(checkpoint)
        
        self.model.to(device)
        self.model.eval()
        
        # 图像预处理
        self.transform = transforms.Compose([
            transforms.Resize((512, 512)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], 
                               std=[0.229, 0.224, 0.225])
        ])
    
    def preprocess_image(self, image):
        """预处理输入图像"""
        if isinstance(image, np.ndarray):
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        
        # 保存原始尺寸用于后处理
        self.orig_size = image.size
        
        # 应用变换
        input_tensor = self.transform(image).unsqueeze(0)
        return input_tensor.to(self.device)
    
    def postprocess_edges(self, edge_map):
        """后处理边缘图"""
        # 转换为numpy数组
        edge_map = edge_map.squeeze().cpu().numpy()
        
        # 调整回原始尺寸
        edge_map = cv2.resize(edge_map, self.orig_size, interpolation=cv2.INTER_LINEAR)
        
        # 归一化到0-255
        edge_map = (edge_map * 255).astype(np.uint8)
        
        return edge_map
    
    def detect_edges(self, image):
        """检测图像边缘"""
        with torch.no_grad():
            # 预处理
            input_tensor = self.preprocess_image(image)
            
            # 推理
            edge_output = self.model(input_tensor)
            
            # 后处理
            edge_map = self.postprocess_edges(edge_output)
            
            return edge_map
    
    def find_document_contours(self, edge_map, min_area=1000):
        """从边缘图中找到文档轮廓"""
        # 形态学操作清理边缘
        kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
        edge_map = cv2.morphologyEx(edge_map, cv2.MORPH_CLOSE, kernel)
        edge_map = cv2.morphologyEx(edge_map, cv2.MORPH_OPEN, kernel)
        
        # 找到轮廓
        contours, _ = cv2.findContours(edge_map, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # 过滤小轮廓
        valid_contours = [cnt for cnt in contours if cv2.contourArea(cnt) > min_area]
        
        # 按面积排序，选择最大的
        if valid_contours:
            largest_contour = max(valid_contours, key=cv2.contourArea)
            
            # 多边形近似
            epsilon = 0.02 * cv2.arcLength(largest_contour, True)
            approx = cv2.approxPolyDP(largest_contour, epsilon, True)
            
            # 如果近似为四边形，认为是文档
            if len(approx) == 4:
                return approx.reshape(4, 2)
        
        return None
    
    def detect_document(self, image):
        """完整的文档检测流程"""
        # 检测边缘
        edge_map = self.detect_edges(image)
        
        # 找到文档轮廓
        corners = self.find_document_contours(edge_map)
        
        return edge_map, corners

# 模型转换为ONNX格式（用于Android）
def convert_to_onnx(model_path, output_path):
    """将PyTorch模型转换为ONNX格式"""
    device = torch.device('cpu')
    model = PiDiNet()
    
    checkpoint = torch.load(model_path, map_location=device)
    if 'state_dict' in checkpoint:
        model.load_state_dict(checkpoint['state_dict'])
    else:
        model.load_state_dict(checkpoint)
    
    model.eval()
    
    # 创建示例输入
    dummy_input = torch.randn(1, 3, 512, 512)
    
    # 导出ONNX
    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        export_params=True,
        opset_version=11,
        do_constant_folding=True,
        input_names=['input'],
        output_names=['output'],
        dynamic_axes={'input': {0: 'batch_size'},
                     'output': {0: 'batch_size'}}
    )
    
    print(f"模型已导出为: {output_path}")

if __name__ == "__main__":
    # 示例使用
    detector = DocumentEdgeDetector("table5_pidinet-tiny-l.pth")
    
    # 转换模型为ONNX
    convert_to_onnx("table5_pidinet-tiny-l.pth", "pidinet_mobile.onnx") 